"""
PostgreSQL数据库连接和会话管理
"""
from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    AsyncEngine,
    create_async_engine,
    async_sessionmaker,
)
from sqlalchemy import text
from sqlalchemy.orm import declarative_base
from sqlalchemy.pool import NullPool, QueuePool
from loguru import logger

from app.config import settings


# 创建Base类，所有模型都继承这个类
Base = declarative_base()


class DatabaseManager:
    """数据库管理器"""

    def __init__(self):
        self.engine: AsyncEngine | None = None
        self.async_session_maker: async_sessionmaker | None = None

    def init(self):
        """初始化数据库连接"""
        try:
            # 构建数据库URL
            database_url = (
                f"postgresql+asyncpg://{settings.POSTGRES_USER}:{settings.POSTGRES_PASSWORD}"
                f"@{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}"
            )

            # 根据环境选择连接池策略
            # 注意：异步引擎会自动使用AsyncAdaptedQueuePool，不需要手动指定
            engine_kwargs = {
                "echo": settings.DEBUG,  # 开发环境打印SQL语句
                "pool_size": settings.DB_POOL_SIZE,
                "max_overflow": settings.DB_MAX_OVERFLOW,
                "pool_timeout": settings.DB_POOL_TIMEOUT,
                "pool_recycle": settings.DB_POOL_RECYCLE,
                "pool_pre_ping": True,  # 检查连接是否有效
            }
            
            # 测试环境使用NullPool
            if settings.ENV == "testing":
                engine_kwargs["poolclass"] = NullPool
                logger.info("使用NullPool连接池（测试环境）")
            else:
                logger.info("使用默认异步连接池（AsyncAdaptedQueuePool）")

            # 创建异步引擎
            self.engine = create_async_engine(database_url, **engine_kwargs)

            # 创建会话工厂
            self.async_session_maker = async_sessionmaker(
                self.engine,
                class_=AsyncSession,
                expire_on_commit=False,  # 提交后不过期对象
                autocommit=False,
                autoflush=False,
            )

            logger.info(
                f"数据库连接初始化成功: {settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}"
            )

        except Exception as e:
            logger.error(f"数据库连接初始化失败: {e}")
            raise

    async def close(self):
        """关闭数据库连接"""
        if self.engine:
            await self.engine.dispose()
            logger.info("数据库连接已关闭")

    async def create_tables(self):
        """创建所有表（开发环境使用，生产环境使用Alembic）"""
        if not self.engine:
            raise RuntimeError("数据库引擎未初始化")

        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
            logger.info("数据库表创建成功")

    async def drop_tables(self):
        """删除所有表（谨慎使用！）"""
        if not self.engine:
            raise RuntimeError("数据库引擎未初始化")

        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.drop_all)
            logger.warning("数据库表已删除")

    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        获取数据库会话（用于依赖注入）
        
        使用示例：
        ```python
        async def get_user(user_id: int, db: AsyncSession = Depends(get_db)):
            result = await db.execute(select(User).where(User.id == user_id))
            return result.scalar_one_or_none()
        ```
        """
        if not self.async_session_maker:
            raise RuntimeError("数据库会话工厂未初始化")

        async with self.async_session_maker() as session:
            try:
                yield session
                await session.commit()
            except Exception as e:
                await session.rollback()
                logger.error(f"数据库会话错误: {e}")
                raise
            finally:
                await session.close()


# 全局数据库管理器实例
db_manager = DatabaseManager()


# 依赖注入函数
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """
    FastAPI依赖注入函数，用于获取数据库会话
    
    使用示例：
    ```python
    @router.get("/users/{user_id}")
    async def read_user(user_id: int, db: AsyncSession = Depends(get_db)):
        # 使用db进行数据库操作
        pass
    ```
    """
    async for session in db_manager.get_session():
        yield session


# 事务装饰器
def transactional(func):
    """
    事务装饰器，自动处理事务提交和回滚
    
    使用示例：
    ```python
    @transactional
    async def create_user_with_profile(db: AsyncSession, user_data: dict):
        user = User(**user_data)
        db.add(user)
        await db.flush()  # 获取user.id
        
        profile = Profile(user_id=user.id, ...)
        db.add(profile)
        # 装饰器会自动commit或rollback
    ```
    """
    async def wrapper(*args, **kwargs):
        # 查找db参数
        db = None
        for arg in args:
            if isinstance(arg, AsyncSession):
                db = arg
                break
        
        if db is None:
            db = kwargs.get('db')
        
        if db is None:
            raise ValueError("未找到AsyncSession参数")

        try:
            result = await func(*args, **kwargs)
            await db.commit()
            return result
        except Exception as e:
            await db.rollback()
            logger.error(f"事务执行失败: {e}")
            raise

    return wrapper


# 数据库健康检查
async def check_database_health() -> bool:
    """
    检查数据库连接是否正常
    
    返回:
        bool: True表示连接正常，False表示连接异常
    """
    try:
        if not db_manager.engine:
            return False

        async with db_manager.engine.connect() as conn:
            await conn.execute(text("SELECT 1"))
        return True
    except Exception as e:
        logger.error(f"数据库健康检查失败: {e}")
        return False


# 数据库初始化函数（在应用启动时调用）
async def init_database():
    """初始化数据库连接"""
    db_manager.init()
    
    # 开发环境自动创建表
    if settings.ENV == "development" and settings.DEBUG:
        logger.info("开发环境：自动创建数据库表")
        await db_manager.create_tables()


# 数据库关闭函数（在应用关闭时调用）
async def close_database():
    """关闭数据库连接"""
    await db_manager.close()

